import cv2
from glob import glob
from datasets.vicon.camera2ob import read_truth_Rt
import numpy as np
from scipy.spatial.transform import Rotation as R


# 刚体1,5个点 采集了位姿

point_ob1 = np.load("/media/liyuke/dataset/vicon/ob1.npy")
point_ob2 = np.load("/media/liyuke/dataset/vicon/ob2.npy")
point_ob3 = np.load("/media/liyuke/dataset/vicon/ob3.npy")
point_ob4 = np.load("/media/liyuke/dataset/vicon/ob4.npy")
point = np.concatenate((point_ob1,point_ob2,point_ob3,point_ob4),axis=0)
K = np.load("/media/liyuke/dataset/vicon4/intrinsic.npy")



if __name__ == "__main__":
    import random
    img_path = "/media/liyuke/dataset/vicon4/seq1/"
    img_files = glob(img_path + '*.jpg')
    img_files = glob(img_path + '*.jpg')
    label_files = list(
        x.strip().replace('.jpg', '_pose.txt') for x in img_files)
    for i in range(len(img_files)):
        img = cv2.imread(img_files[i])
        image_h,image_w = img.shape[:2]
        rotation_matrix, translation_matrix = read_truth_Rt(label_files[i])
        p_camera = np.dot(rotation_matrix, point.T) + translation_matrix
        p_img = np.dot(K, p_camera)
        p_img = (p_img[:2, :] / p_img[2,]).T
        # 计算外接矩形
        min_x = int(min(p_img,key=lambda p:p[0])[0])
        min_x = min(max(min_x, 0), image_w)
        max_x = int(max(p_img,key=lambda p:p[0])[0])
        max_x = min(max(max_x, 0), image_w)
        min_y = int(min(p_img,key=lambda p:p[1])[1])
        min_y = min(max(min_y, 0), image_h)
        max_y = int(max(p_img,key=lambda p:p[1])[1])
        max_y = min(max(max_y, 0), image_h)
        x_center = (min_x+max_x)/2/image_w
        y_center = (min_y+max_y)/2/image_h
        w = (max_x-min_x)/image_w
        h = (max_y-min_y)/image_h
        data = [str(0), str(x_center), str(y_center), str(w), str(h)]
        # with open(img_files[i].replace("jpg", "txt"), "w") as f:
        #     f.write(" ".join(data))
        for i in range(20):
            # cv2.circle(img, (int(p_img[0][i]), int(p_img[1][i])), 3, (0, 0, 255), -1)
            cv2.rectangle(img,(min_x,min_y),(max_x,max_y),(0,255,0),2)
        cv2.imshow("A", img)
        cv2.waitKey()
